import json
from datetime import datetime, time

import chardet
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from adjustText import adjust_text
from matplotlib import ticker
from matplotlib.lines import Line2D
from matplotlib.patches import FancyArrowPatch
from matplotlib.ticker import MaxNLocator

# 全局配置
mpl.rcParams.update({
    'font.family': 'Microsoft YaHei',
    'axes.unicode_minus': False,
    'agg.path.chunksize': 20000
})

# 文件路径
kline_trade_file = 'D:\\work\\code\\clark\\gitee\\big_a\\datas\\backtest\\sh513180\\20250714-202507181500_chart.txt'

def detect_encoding(file_path):
    with open(file_path, 'rb') as f:
        return chardet.detect(f.read(10000))['encoding']

def load_kline_data(file_path):
    kline_data = []
    encoding = detect_encoding(file_path)
    with open(file_path, 'r', encoding=encoding, errors='replace') as file:
        for line in file:
            if line.strip().startswith('{"add":'):
                try:
                    kline_data.append(json.loads(line))
                except:
                    continue
    return kline_data

def load_trade_data(file_path):
    trade_data = []
    encoding = detect_encoding(file_path)
    with open(file_path, 'r', encoding=encoding, errors='replace') as file:
        for line in file:
            if line.strip().startswith('{"cash":'):
                try:
                    trade_data.append(json.loads(line))
                except:
                    continue
    return trade_data

# 加载数据
kline_data = load_kline_data(kline_trade_file)
trade_data = load_trade_data(kline_trade_file)

# 过滤交易时间段数据
def is_trading_time(dt):
    t = dt.time()
    return (time(9, 30) <= t <= time(11, 30)) or (time(13, 0) <= t <= time(15, 0))

# 准备K线数据
raw_dates = [datetime.strptime(k['date'], '%Y-%m-%d %H:%M:%S') for k in kline_data]
trading_mask = [is_trading_time(dt) for dt in raw_dates]

dates = [dt for dt, mask in zip(raw_dates, trading_mask) if mask]
opens = [k['open'] for k, mask in zip(kline_data, trading_mask) if mask]
highs = [k['high'] for k, mask in zip(kline_data, trading_mask) if mask]
lows = [k['low'] for k, mask in zip(kline_data, trading_mask) if mask]
closes = [k['close'] for k, mask in zip(kline_data, trading_mask) if mask]
volumes = [k['volume'] for k, mask in zip(kline_data, trading_mask) if mask]

# === 创建数值索引代替时间轴 ===
x_index = np.arange(len(dates))  # 连续的数值索引

# 创建图表
fig = plt.figure(figsize=(18, 14), layout="constrained")
gs = fig.add_gridspec(nrows=3, ncols=1, height_ratios=[4, 1, 1], hspace=0.05)

# 1. K线图
ax1 = fig.add_subplot(gs[0])
ax1.set_title(f'资产价格走势与交易点', fontsize=14, fontweight='bold')
ax1.set_ylabel('价格', fontsize=12)
ax1.grid(True, linestyle='--', alpha=0.7)

# 动态计算K线宽度
width = 0.6  # 固定宽度更稳定

# 绘制K线（使用数值索引）
for i in range(len(x_index)):
    color = 'red' if closes[i] >= opens[i] else 'green'
    ax1.vlines(x_index[i], lows[i], highs[i], color=color, linewidth=0.8)
    ax1.bar(x_index[i],
            height=abs(closes[i]-opens[i]),
            bottom=min(opens[i], closes[i]),
            width=width,
            color=color,
            edgecolor=color)

# 添加交易时段分割线（使用数值索引）
for i, dt in enumerate(dates):
    if dt.time() == time(13, 0):
        ax1.axvline(x_index[i], color='red', linestyle='--', alpha=0.4, linewidth=0.8)

# 2. 成交量（使用数值索引）
ax2 = fig.add_subplot(gs[1], sharex=ax1)
ax2.bar(x_index, volumes, width=width*0.8,
        color=np.where(np.array(closes) >= np.array(opens), 'red', 'green'))
ax2.set_ylabel('成交量', fontsize=12)
ax2.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{x/1e6:.1f}M'))
ax2.grid(True, linestyle='--', alpha=0.7)

# 3. 资产曲线
ax3 = fig.add_subplot(gs[2], sharex=ax1)
trade_dates = [datetime.strptime(t['date'], '%Y-%m-%d %H:%M:%S') for t in trade_data]
assets = [t['totalAsset'] for t in trade_data]

# 映射交易时间到K线索引位置
trade_indices = []
for td in trade_dates:
    min_idx = min(range(len(dates)), key=lambda i: abs(dates[i] - td))
    trade_indices.append(x_index[min_idx])

# 绘制资产曲线（使用数值索引）
ax3.plot(trade_indices, assets, 'b-', linewidth=2, alpha=0.8)
ax3.fill_between(trade_indices, min(assets), assets, color='blue', alpha=0.1)
ax3.set_ylabel('总资产', fontsize=12)
ax3.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'¥{x/1000:.0f}K'))
ax3.grid(True, linestyle='--', alpha=0.7)

# 标记起始点和终点（使用数值索引）
ax3.scatter(trade_indices[0], assets[0], s=80, color='green', edgecolor='white', zorder=5,
            label=f'起始: ¥{assets[0]:.3f}')
ax3.scatter(trade_indices[-1], assets[-1], s=80, color='red', edgecolor='white', zorder=5,
            label=f'结束: ¥{assets[-1]:.3f}')

# 标记资产峰值（使用数值索引）
peak_idx = np.argmax(assets)
ax3.scatter(trade_indices[peak_idx], assets[peak_idx], s=100, color='gold', edgecolor='darkorange',
            marker='*', zorder=6, label=f'峰值: ¥{assets[peak_idx]/1000:.0f}K')
ax3.legend(loc='upper left', fontsize=9)

# === 优化刻度设置：均匀分布刻度，避免重叠 ===
if dates:
    # 使用MaxNLocator确保刻度数量在10个以内
    ax1.xaxis.set_major_locator(MaxNLocator(nbins=10, integer=True))

    # 获取刻度位置
    tick_positions = ax1.get_xticks()
    # 过滤出在有效范围内的刻度位置
    valid_indices = [int(pos) for pos in tick_positions if 0 <= pos < len(dates)]

    # 生成对应的日期标签
    tick_labels = [dates[i].strftime('%m-%d %H:%M') for i in valid_indices]

    # 应用刻度设置
    for ax in [ax1, ax2, ax3]:
        ax.set_xticks(valid_indices)
        ax.set_xticklabels(tick_labels, rotation=45, ha='right', fontsize=9)  # 调整字体大小提高可读性
else:
    # 如果没有数据，设置空刻度
    for ax in [ax1, ax2, ax3]:
        ax.set_xticks([])
        ax.set_xticklabels([])

# 识别买卖点
buy_points, sell_points = [], []
for trade in trade_data:
    dt = datetime.strptime(trade['date'], '%Y-%m-%d %H:%M:%S')
    vt = trade['date'][11:16]
    price = trade['price']
    operation = trade['operation']
    volume = trade.get('volume', trade.get('position', 0))

    min_idx = min(range(len(dates)), key=lambda i: abs(dates[i] - dt))
    if "买入" == operation:
        buy_points.append((x_index[min_idx], price, volume, min_idx, vt))
    elif "卖出" == operation:
        sell_points.append((x_index[min_idx], price, volume, min_idx, vt))

# 绘制买卖点标记（使用数值索引）
texts = []
for point in buy_points:
    x_pos, price, volume, idx, vt = point
    t = ax1.text(x_pos, price, f'B {price:.3f}\n+{volume} {vt}',
                 fontsize=9, color='red', ha='center', va='top',
                 bbox=dict(boxstyle="round,pad=0.2", fc='white', ec='red', alpha=0.8))
    texts.append(t)
    # 添加箭头
    ax1.add_patch(FancyArrowPatch(
        (x_pos, price*0.98), (x_pos, price),
        arrowstyle='->', color='red', lw=1, alpha=0.7))

for point in sell_points:
    x_pos, price, volume, idx, vt = point
    t = ax1.text(x_pos, price, f'S {price:.3f}\n-{volume} {vt}',
                 fontsize=9, color='green', ha='center', va='bottom',
                 bbox=dict(boxstyle="round,pad=0.2", fc='white', ec='green', alpha=0.8))
    texts.append(t)
    # 添加箭头
    ax1.add_patch(FancyArrowPatch(
        (x_pos, price*1.02), (x_pos, price),
        arrowstyle='->', color='green', lw=1, alpha=0.7))

# 智能调整标注
adjust_text(texts,
            arrowprops=dict(arrowstyle='->', color='gray', lw=0.8, alpha=0.7, shrinkA=10),
            expand_points=(1.5, 1.5),
            expand_text=(1.2, 1.2),
            force_points=0.5,
            ax=ax1)

# 添加图例
legend_elements = [
    Line2D([0], [0], color='red', lw=4, label='阳线'),
    Line2D([0], [0], color='green', lw=4, label='阴线'),
    Line2D([0], [0], marker='o', color='red', label='买入点', markersize=8, linestyle='None'),
    Line2D([0], [0], marker='o', color='green', label='卖出点', markersize=8, linestyle='None'),
    Line2D([0], [0], color='red', linestyle='--', alpha=0.4, label='交易时段分割')
]
ax1.legend(handles=legend_elements, loc='upper left', fontsize=9)

# 调整布局
plt.subplots_adjust(left=0.06, right=0.95, top=0.95, bottom=0.12, hspace=0.15)
plt.savefig('专业K线分析图.png', dpi=150, bbox_inches='tight')
plt.show()